"""
土壤水分数据处理
nc文件转tif
"""
import os

import netCDF4 as nc
import numpy as np
from osgeo import gdal, osr, ogr


def nc2tif(file, output_folder):
    pre_data = nc.Dataset(file)  # 利用.Dataset()读取nc数据
    # print(pre_data.variables.keys())
    Lat_data = pre_data.variables['lat'][:]
    Lon_data = pre_data.variables['lon'][:]
    pre_arr = np.squeeze(pre_data.variables['sm'])  # 属性变量名
    # print(pre_arr)
    # print(pre_arr.shape)
    # 影像的左上角&右下角坐标
    Lonmin, Latmax, Lonmax, Latmin = [Lon_data.min(), Lat_data.max(), Lon_data.max(), Lat_data.min()]
    # Lonmin, Latmax, Lonmax, Latmin = [-180, 84.656, 180, -84.656]
    # 分辨率计算
    Num_lat = len(Lat_data)
    Num_lon = len(Lon_data)
    Lat_res = (Latmax - Latmin) / (float(Num_lat) - 1)
    Lon_res = (Lonmax - Lonmin) / (float(Num_lon) - 1)
    # print("分辨率：", Lat_res, Lon_res)
    # 创建tif文件
    driver = gdal.GetDriverByName('GTiff')
    out_tif_name = output_folder + "/" + os.path.basename(file).split('-')[5][0:8] + '.tif'
    out_tif = driver.Create(out_tif_name, Num_lon, Num_lat, 1, gdal.GDT_Float32)
    # 设置影像的显示范围
    # Lat_re前需要添加负号
    geotransform = (Lonmin, Lon_res, 0.0, Latmax, 0.0, -Lat_res)
    out_tif.SetGeoTransform(geotransform)
    # 定义投影
    prj = osr.SpatialReference()
    prj.ImportFromEPSG(4326)
    out_tif.SetProjection(prj.ExportToWkt())
    # 数据导出
    # print(pre_arr.shape)
    # print(out_tif.GetRasterBand(1).ReadAsArray().shape)
    # pre_arr = np.transpose(pre_arr)
    out_tif.GetRasterBand(1).WriteArray(pre_arr)  # 将数据写入内存
    out_tif.GetRasterBand(1).SetNoDataValue(-9999)
    out_tif.FlushCache()  # 将数据写入到硬盘
    del out_tif  # 关闭tif文件
    print("转换完成：{}".format(file))


def main(input_dir, output_dir):
    file_ls = [os.path.join(input_dir, f) for f in os.listdir(input_dir) if f.endswith(".nc")]
    # print(file_ls)

    for i, file in enumerate(file_ls):
        nc2tif(file, output_dir)


def transAndMean(input1_dir, input2_dir, output_dir):
    file1_ls = [os.path.join(input1_dir, f) for f in os.listdir(input1_dir) if f.endswith(".nc4")]
    file2_ls = [os.path.join(input2_dir, f) for f in os.listdir(input2_dir) if f.endswith(".nc4")]

    for file1, file2 in zip(file1_ls, file2_ls):
        pre_data1 = nc.Dataset(file1)  # 利用.Dataset()读取nc数据
        pre_data2 = nc.Dataset(file2)  # 利用.Dataset()读取nc数据
        print(pre_data1.variables.keys())
        Lat_data = pre_data1.variables['lat'][:]
        Lon_data = pre_data1.variables['lon'][:]
        pre_arr1 = np.squeeze(pre_data1.variables['sm'])  # 属性变量名
        pre_arr2 = np.squeeze(pre_data2.variables['sm'])  # 属性变量名
        pre_arr1[np.isnan(pre_arr1)] = 0
        pre_arr2[np.isnan(pre_arr2)] = 0
        pre_arr = (pre_arr1 + pre_arr2) / 2

        pre_arr[pre_arr == 0] = -9999
        # print(pre_arr)
        # print(pre_arr.shape)
        # 影像的左上角&右下角坐标
        Lonmin, Latmax, Lonmax, Latmin = [Lon_data.min(), Lat_data.max(), Lon_data.max(), Lat_data.min()]
        # Lonmin, Latmax, Lonmax, Latmin = [-180, 84.656, 180, -84.656]
        # 分辨率计算
        Num_lat = len(Lat_data)
        Num_lon = len(Lon_data)
        Lat_res = (Latmax - Latmin) / (float(Num_lat) - 1)
        Lon_res = (Lonmax - Lonmin) / (float(Num_lon) - 1)
        # print("分辨率：", Lat_res, Lon_res)
        # 创建tif文件
        driver = gdal.GetDriverByName('GTiff')
        out_tif_name = output_dir + "/" + os.path.basename(file1).split("_")[-2] + '.tif'
        out_tif = driver.Create(out_tif_name, Num_lon, Num_lat, 1, gdal.GDT_Float32)
        # 设置影像的显示范围
        # Lat_re前需要添加负号
        geotransform = (Lonmin, Lon_res, 0.0, Latmax, 0.0, -Lat_res)
        out_tif.SetGeoTransform(geotransform)
        # 定义投影
        prj = osr.SpatialReference()
        prj.ImportFromEPSG(4326)
        out_tif.SetProjection(prj.ExportToWkt())
        # 数据导出
        # print(pre_arr.shape)
        # print(out_tif.GetRasterBand(1).ReadAsArray().shape)
        pre_arr = np.transpose(pre_arr)
        out_tif.GetRasterBand(1).WriteArray(pre_arr)  # 将数据写入内存
        out_tif.GetRasterBand(1).SetNoDataValue(-9999)
        out_tif.FlushCache()  # 将数据写入到硬盘
        del out_tif  # 关闭tif文件
        print("转换完成：{}".format(out_tif_name))


def cmorph_nc2tif(file, output_folder):
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)
    pre_data = nc.Dataset(file)  # 利用.Dataset()读取nc数据
    # print(pre_data.variables.keys())
    # print(pre_data.variables['lat'][:].shape)
    Lat_data = pre_data.variables['lat'][:]
    Lon_data = pre_data.variables['lon'][:]
    pre_arr = np.squeeze(pre_data.variables['cmorph'])  # 属性变量名
    # print(pre_arr)
    # print(pre_arr.shape)
    # 影像的左上角&右下角坐标
    Lonmin, Latmax, Lonmax, Latmin = [Lon_data.min(), Lat_data.max(), Lon_data.max(), Lat_data.min()]
    # Lonmin, Latmax, Lonmax, Latmin = [-180, 84.656, 180, -84.656]
    # 分辨率计算
    Num_lat = len(Lat_data)
    Num_lon = len(Lon_data)
    Lat_res = (Latmax - Latmin) / (float(Num_lat) - 1)
    Lon_res = (Lonmax - Lonmin) / (float(Num_lon) - 1)
    # print("分辨率：", Lat_res, Lon_res)
    # 创建tif文件
    driver = gdal.GetDriverByName('GTiff')
    out_tif_name = output_folder + "/" + os.path.basename(file).split('_')[-1][0:8] + '.tif'
    out_tif = driver.Create(out_tif_name, Num_lon, Num_lat, 1, gdal.GDT_Float32)
    # 设置影像的显示范围
    # Lat_re前需要添加负号
    geotransform = (Lonmin, Lon_res, 0.0, Latmax, 0.0, -Lat_res)
    out_tif.SetGeoTransform(geotransform)
    # 定义投影
    prj = osr.SpatialReference()
    prj.ImportFromEPSG(4326)
    out_tif.SetProjection(prj.ExportToWkt())
    # 数据导出
    # print(pre_arr.shape)
    # print(out_tif.GetRasterBand(1).ReadAsArray().shape)
    # pre_arr = np.transpose(pre_arr)
    out_tif.GetRasterBand(1).WriteArray(pre_arr)  # 将数据写入内存
    out_tif.GetRasterBand(1).SetNoDataValue(-9999)
    out_tif.FlushCache()  # 将数据写入到硬盘
    del out_tif  # 关闭tif文件
    print("转换完成：{}".format(file))

if __name__ == '__main__':
    # CCI_path = r"G:\test\SMAP_NEW\ESACCI\2015\ESACCI-SOILMOISTURE-L3S-SSMV-COMBINED-20150101000000-fv08.1.nc"
    # nc2tif(CCI_path, "")
    # input1_dir = r"G:\test\SMAP_NEW\MCCA\NC\Des_2015\2015"
    # input2_dir = r"G:\test\SMAP_NEW\MCCA\NC\Asc_2015\2015"
    # output_dir = r"G:\test\SMAP_NEW\ESACCI\TIF\2014-2017"
    # transAndMean(input1_dir, input2_dir, output_dir)

    # input_dir = r"G:\test\SMAP_NEW\ESACCI\2014"
    # main(input_dir, output_dir)

    input_dir = r"G:\test\SMAP_NEW\CMORPH"
    output_dir = r"G:\test\SMAP_NEW\CMORPH_tif"
    # cmorph_path = r"G:\test\SMAP_NEW\CMORPH\CMORPH_V1.0_ADJ_0.25deg-DLY_00Z_20150103.nc"
    file_ls = [os.path.join(input_dir, f) for f in os.listdir(input_dir) if f.endswith(".nc")]
    # print(file_ls)

    for i, file in enumerate(file_ls):
        cmorph_nc2tif(file, output_dir)
